import numpy
import cv2


# 4.3 从视差图得到掩膜
def create_median_mask(disparity_map, valid_depth_mask, rect=None):
    """Return a mask selecting the median layer, plus shadows."""
    if rect is not None:
        x, y, w, h = rect
        disparity_map = disparity_map[y:y + h, x:x + w]
        valid_depth_mask = valid_depth_mask[y:y + h, x:x + w]
    median = numpy.median(disparity_map)

    """ 
    numpy.where有三个参数：第一个参数是数组，该数组的元素值为真或假。
    函数会返回同样维度的数组；
    输入数组的元素为真时，输出数组的相应元素为函数的第二个参数；
    否则，输入数组的元素为假时，输出数组相应元素为函数的第三个参数。
    """
    return numpy.where(
        (valid_depth_mask == 0) | (abs(disparity_map - median) < 12),
        1.0,
        0.0
    )


# 4.5 使用普通摄像头进行深度估计
def update():
    # disparity range is turned for 'aloe' image pair.
    stereo.setBlockSize(cv2.getTrackbarPos('window_size', 'disparity'))
    stereo.setUniquenessRatio(cv2.getTrackbarPos('uniquenessRatio', 'disparity'))
    stereo.setSpeckleWindowSize(cv2.getTrackbarPos('speckleRange', 'disparity'))
    stereo.setDisp12MaxDiff(cv2.getTrackbarPos('disp12MaxDiff', 'disparity'))

    print ("computing disparity...")
    disparity = stereo.compute(left_img, right_img).astype(numpy.float32) / 16.0

    cv2.imshow('left', left_img)
    cv2.imshow('disparity', (disparity - min_disparity) / num_disparity)


if __name__ == '__main__':
    window_size = 5
    min_disparity = 16
    num_disparity = 192-min_disparity
    block_size = window_size
    uniqueness_ratio = 1
    speckle_range = 3
    speckle_window_size = 3
    disparity_max_diff = 200
    p1 = 600
    p2 = 2400
    left_img = cv2.imread('image/1.jpg')
    right_img = cv2.imread('image/2.jpg')
    cv2.namedWindow('disparity')
    cv2.createTrackbar('speckleRange', 'disparity', speckle_range, 50, update)
    cv2.createTrackbar('window_size', 'disparity', window_size, 21, update)
    cv2.createTrackbar('speckleWindowSize', 'disparity', speckle_window_size, 200, update)
    cv2.createTrackbar('uniquenessRatio', 'disparity', uniqueness_ratio, 50, update)
    cv2.createTrackbar('disp12MaxDiff', 'disparity', disparity_max_diff, 250, update)
    stereo = cv2.StereoSGBM_create(
        minDisparity=min_disparity,             # 表示可能的最小视差值。它通常为0，但有时候校正算法会移动图像，所以参数也要相应调整。
        numDisparities=num_disparity,           # 表示最大的视差值和最小的视差值之差。这个差值总是大于0，在当前的视线中，这个值必须要能被16整除
        blockSize=block_size,
        uniquenessRatio=uniqueness_ratio,       # 表示由代价函数计算得到的最好（最小）结果值比第二好的值小多少（用百分比表示）才被认为是正确的。通常在5~15之间就可以了。
        speckleRange=speckle_range,             # 表示每个已连接部分的最大视差变化。如果进行斑点过滤，则该参数取正值，函数会自动乘以16,。一般情况下取1或2就足够了。
        speckleWindowSize=speckle_window_size,  # 表示平滑视差区域的最大窗口尺寸，以考虑噪声斑点的无效性。将它设为0就不会进行斑点过滤。否则应取50~200之间的某个值
        disp12MaxDiff=disparity_max_diff,       # 表示在左右视差检查中最大允许的偏差（整数像素单位）.设为非正值将不做检查。
        P1=p1,                                  # 控制视差平滑度的第一个参数
        P2=p2                                   # 控制视差平滑度的第二个参数
    )
    update()
    cv2.waitKey()
    cv2.destroyAllWindows()
